//
// Created by Luhang on 2023/5/25/025.
//

#ifndef INC_0423_AVLTREE_H
#define INC_0423_AVLTREE_H

#endif //INC_0423_AVLTREE_H

#include <iostream>
#include <cassert>
#include <ctime>
using namespace std;

template<class K, class V>
struct AVLTreeNode{
    AVLTreeNode<K, V>* _left;
    AVLTreeNode<K, V>* _right;
    AVLTreeNode<K, V>* _parent;
    pair<K, V> _kv;
    int _bf;// balance factor

    AVLTreeNode(const pair<K, V>& kv)
    :_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    ,_kv(kv)
    ,_bf(0){}
};

template<class K, class V>
class AVLTree{
    typedef AVLTreeNode<K, V> Node;
public:
    bool Insert(const pair<K, V>& kv){
        if (_root == nullptr){
            _root = new Node(kv);
            return true;
        }

        Node* parent = nullptr;
        Node* cur = _root;

        //先定位
        while (cur){
            if (cur->_kv.first < kv.first){
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first){
                parent = cur;
                cur = cur->_left;
            }
            else
                return false;
        }

        //链接
        cur = new Node(kv);
        if (parent->_kv.first < kv.first)
            parent->_right = cur;
        else
            parent->_left = cur;
        cur->_parent = parent;

        //修正并更新平衡因子
        while (parent){
            if (cur == parent->_right)
                ++parent->_bf;
            else
                --parent->_bf;

            if (parent->_bf == 1 || parent->_bf == -1){
                parent = parent->_parent;
                cur = cur->_parent;
            }
            else if (parent->_bf == 0)
                break;
            else if (parent->_bf == 2 || parent->_bf == -2){
                //旋转
                if (parent->_bf == 2 && cur->_bf == 1)
                    RotateL(parent);
                else if (parent->_bf == -2 && cur->_bf == -1)
                    RotateR(parent);
                else if (parent->_bf == 2 && cur->_bf == -1)
                    RotateRL(parent);
                else if (parent->_bf == -2 && cur->_bf == 1)
                    RotateLR(parent);
                else
                    assert(false);
                break;
            } else
                assert(false);
        }
        return true;
    }

    void InOrder(){
        _InOrder(_root);
        cout << endl;
    }

    bool IsBalance(){
        return _IsBalance(_root);
    }

    int Height(){
        return _Height(_root);
    }

private:
    void _InOrder(Node* root){
        if (root == nullptr)
            return;
        _InOrder(root->_left);
        cout << root->_kv.first << " ";
        _InOrder(root->_right);
    }

    bool _IsBalance(Node* root){
        if (root == nullptr)
            return true;

        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);

        if (rightHeight - leftHeight != root->_bf){
            cout << root->_kv.first << "balance factor wrong" << endl;
            return false;
        }

        return abs(leftHeight - rightHeight) < 2 && _IsBalance(root->_right) && _IsBalance(root->_left);
    }

    int _Height(Node* root){
        if (root == nullptr)
            return 0;
        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);

        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

    void RotateL(Node* parent){
        Node* subR = parent->_right;
        Node* subRL = subR->_left;

        parent->_right = subRL;
        if (subRL)
            subRL->_parent = parent;

        Node* pParent = parent->_parent;

        subR->_left = parent;
        parent->_parent = subR;

        if (pParent == nullptr){
            _root = subR;
            subR->_parent = nullptr;
        } else{
            if (pParent->_right == parent)
                pParent->_right = subR;
            else
                pParent->_left = subR;

            subR->_parent = pParent;
        }

        parent->_bf = subR->_bf = 0;
    }

    void RotateR(Node* parent){
        Node* subL = parent->_left;
        Node* subLR = subL->_right;

        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;

        Node* pParent = parent->_parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root){
            _root = subL;
            subL->_parent = nullptr;
        } else{
            if (pParent->_right == parent)
                pParent->_right = subL;
            else
                pParent->_left = subL;
            subL->_parent = pParent;
        }

        subL->_bf = parent->_bf = 0;
    }

    void RotateLR(Node* parent){
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        int bf = subLR->_bf;

        RotateL(subL);
        RotateR(parent);

        if (bf == 1){
            parent->_bf = 0;
            subL->_bf = -1;
            subLR->_bf = 0;
        }
        else if (bf == -1){
            parent->_bf = 1;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else if (bf == 0){
            parent->_bf = 0;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else
            assert(false);
    }

    void RotateRL(Node* parent){
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        int bf = subRL->_bf;

        RotateR(subR);
        RotateL(parent);

        if (bf == 1){
            parent->_bf = -1;
            subR->_bf = 0;
            subRL->_bf = 0;
        }
        else if (bf == -1){
            parent->_bf = 0;
            subR->_bf = 1;
            subRL->_bf = 0;
        }
        else if (bf == 0){
            parent->_bf = 0;
            subR->_bf = 0;
            subRL->_bf = 0;
        }
        else
            assert(false);
    }

    Node* _root = nullptr;
};

void Test_AVLTree1(){
    int a[] = { 4, 2, 6, 1, 3, 5, 15, 7, 16, 14 };
    AVLTree<int, int> t1;
    for (auto e : a)
    {
        t1.Insert(make_pair(e, e));
        cout << e <<" Insert: "<<t1.IsBalance() << endl;
    }

    t1.InOrder();
    cout << t1.IsBalance() << endl;
}

void Test_AVLTree2()
{
    srand(time(0));
    const size_t N = 10000000;
    AVLTree<int, int> t;
    for (size_t i = 0; i < N; ++i)
    {
        size_t x = rand() + i;
        t.Insert(make_pair(x, x));
        //cout << t.IsBalance() << endl;
    }

    //t.InOrder();

    cout << t.IsBalance() << endl;
    cout << t.Height() << endl;
}